from sklearn.metrics import accuracy_score
from dataLoader import load_dataset
from transformers import TrainingArguments, Trainer
import numpy as np
from models import MyNeuralNetwork

max_len = 50
batch_size = 32
embed_size = 100
hidden_size = 256
output_size = 11
num_epochs = 30

train_dataset, dev_dataset, vocab = load_dataset(max_len)

def compute_metrics(eval_pred):
    logits, labels = eval_pred
    predictions = np.argmax(logits, axis=-1)
    accuracy = accuracy_score(labels, predictions)
    return {"accuracy": accuracy}

# 实例化你的模型
model = MyNeuralNetwork(
    embed_size=embed_size,
    hidden_size=hidden_size,
    output_size=output_size,
    vocab_size=len(vocab)
)

# 定义训练参数
training_args = TrainingArguments(
    output_dir="./results",
    num_train_epochs=num_epochs,
    per_device_train_batch_size=batch_size,
    per_device_eval_batch_size=batch_size,
    learning_rate=0.001,
    evaluation_strategy="epoch",       # 每轮训练后评估
    save_strategy="epoch",             # 每轮训练后保存
    save_total_limit=3,                # 最多保留 3 个 checkpoint
    load_best_model_at_end=True,       # 训练完加载最优模型
    metric_for_best_model="accuracy",  # 选择最优模型的指标
    greater_is_better=True,            # 指标越大越好（accuracy 越高越好）
    logging_dir='./logs',
    logging_steps=100,
    report_to="none",                  # 禁用集成报告工具
)

# 实例化 Trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=dev_dataset,
    compute_metrics=compute_metrics,
)

# 开始训练
trainer.train()
